#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
增强的旋转校正模块 - 专门用于参考图像自动摆正
"""

import cv2
import numpy as np
import numpy.typing as npt
from typing import Tuple, List, Optional
import logging
import math

class EnhancedRotationCorrection:
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        
    def auto_correct_reference_image(self, image: npt.NDArray) -> Tuple[npt.NDArray, float, dict]:
        """
        自动校正参考图像的旋转角度
        
        Args:
            image: 输入图像
            
        Returns:
            (校正后的图像, 旋转角度, 检测信息)
        """
        self.logger.info("开始自动校正参考图像旋转...")
        
        # 方法1: 基于键盘边缘检测
        result1 = self._correct_by_keyboard_edges(image)
        
        # 方法2: 基于字符方向检测
        result2 = self._correct_by_text_orientation(image)
        
        # 方法3: 基于按键排列检测
        result3 = self._correct_by_key_alignment(image)
        
        # 选择最佳结果
        best_result = self._select_best_correction([result1, result2, result3])
        
        if best_result:
            corrected_image, angle, info = best_result
            self.logger.info(f"参考图像已校正旋转 {angle:.1f}° (方法: {info['method']})")
            return corrected_image, angle, info
        else:
            self.logger.warning("无法自动校正图像旋转，保持原图")
            return image, 0.0, {"method": "none", "confidence": 0.0}
    
    def _correct_by_keyboard_edges(self, image: npt.NDArray) -> Optional[Tuple[npt.NDArray, float, dict]]:
        """基于键盘边缘检测的旋转校正"""
        try:
            # 转换为灰度图像
            if len(image.shape) == 3:
                gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            else:
                gray = image.copy()
            
            # 高斯模糊减少噪声
            blurred = cv2.GaussianBlur(gray, (5, 5), 0)
            
            # 边缘检测
            edges = cv2.Canny(blurred, 50, 150, apertureSize=3)
            
            # 形态学操作连接断开的边缘
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
            edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
            
            # 使用霍夫变换检测直线
            lines = cv2.HoughLines(edges, 1, np.pi/180, threshold=100)
            
            if lines is not None and len(lines) > 10:
                # 分析直线角度
                angles = []
                for line in lines:
                    rho, theta = line[0]
                    angle = theta * 180 / np.pi
                    # 标准化角度到 [-45, 45] 范围
                    if angle > 45:
                        angle -= 90
                    elif angle < -45:
                        angle += 90
                    angles.append(angle)
                
                # 计算主要角度
                angles = np.array(angles)
                # 使用中位数更稳定
                main_angle = np.median(angles)
                
                # 计算角度分布的集中度
                angle_std = np.std(angles)
                confidence = max(0, 1.0 - angle_std / 45.0)  # 角度越集中，置信度越高
                
                if abs(main_angle) > 1.0 and confidence > 0.3:  # 至少1度且有一定置信度
                    corrected_image = self._rotate_image(image, -main_angle)
                    return corrected_image, main_angle, {
                        "method": "keyboard_edges",
                        "confidence": confidence,
                        "angle_std": angle_std,
                        "line_count": len(lines)
                    }
            
            return None
            
        except Exception as e:
            self.logger.error(f"键盘边缘检测失败: {e}")
            return None
    
    def _correct_by_text_orientation(self, image: npt.NDArray) -> Optional[Tuple[npt.NDArray, float, dict]]:
        """基于字符方向检测的旋转校正"""
        try:
            # 转换为灰度图像
            if len(image.shape) == 3:
                gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            else:
                gray = image.copy()
            
            # 二值化
            _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
            
            # 查找轮廓
            contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            
            # 筛选可能的字符轮廓
            char_contours = []
            for contour in contours:
                area = cv2.contourArea(contour)
                if 100 < area < 5000:  # 字符大小范围
                    x, y, w, h = cv2.boundingRect(contour)
                    aspect_ratio = w / h
                    if 0.3 < aspect_ratio < 3.0:  # 字符长宽比范围
                        char_contours.append(contour)
            
            if len(char_contours) < 5:  # 至少需要5个字符
                return None
            
            # 分析字符的排列方向
            angles = []
            for contour in char_contours:
                # 使用最小外接矩形
                rect = cv2.minAreaRect(contour)
                angle = rect[2]
                
                # 标准化角度
                if angle < -45:
                    angle += 90
                elif angle > 45:
                    angle -= 90
                
                angles.append(angle)
            
            if angles:
                # 计算主要角度
                angles = np.array(angles)
                main_angle = np.median(angles)
                angle_std = np.std(angles)
                confidence = max(0, 1.0 - angle_std / 30.0)
                
                if abs(main_angle) > 1.0 and confidence > 0.4:
                    corrected_image = self._rotate_image(image, -main_angle)
                    return corrected_image, main_angle, {
                        "method": "text_orientation",
                        "confidence": confidence,
                        "char_count": len(char_contours),
                        "angle_std": angle_std
                    }
            
            return None
            
        except Exception as e:
            self.logger.error(f"字符方向检测失败: {e}")
            return None
    
    def _correct_by_key_alignment(self, image: npt.NDArray) -> Optional[Tuple[npt.NDArray, float, dict]]:
        """基于按键排列检测的旋转校正"""
        try:
            # 转换为灰度图像
            if len(image.shape) == 3:
                gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            else:
                gray = image.copy()
            
            # 预处理
            blurred = cv2.GaussianBlur(gray, (5, 5), 0)
            edges = cv2.Canny(blurred, 50, 150)
            
            # 查找轮廓
            contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            
            # 筛选可能的按键轮廓
            key_contours = []
            for contour in contours:
                area = cv2.contourArea(contour)
                if 500 < area < 20000:  # 按键大小范围
                    x, y, w, h = cv2.boundingRect(contour)
                    aspect_ratio = w / h
                    if 0.5 < aspect_ratio < 2.0:  # 按键长宽比范围
                        key_contours.append(contour)
            
            if len(key_contours) < 10:  # 至少需要10个按键
                return None
            
            # 分析按键的排列方向
            # 计算按键中心点
            centers = []
            for contour in key_contours:
                M = cv2.moments(contour)
                if M["m00"] != 0:
                    cx = int(M["m10"] / M["m00"])
                    cy = int(M["m01"] / M["m00"])
                    centers.append((cx, cy))
            
            if len(centers) < 10:
                return None
            
            # 使用RANSAC拟合直线来检测排列方向
            angles = []
            centers = np.array(centers)
            
            # 尝试不同的角度
            for angle in np.arange(-45, 46, 1):
                # 旋转中心点
                cos_a = np.cos(np.radians(angle))
                sin_a = np.sin(np.radians(angle))
                
                # 计算旋转后的y坐标
                rotated_y = centers[:, 0] * sin_a + centers[:, 1] * cos_a
                
                # 计算y坐标的方差（越小说明排列越整齐）
                y_variance = np.var(rotated_y)
                angles.append((angle, y_variance))
            
            # 选择方差最小的角度
            angles.sort(key=lambda x: x[1])
            best_angle, best_variance = angles[0]
            
            # 计算置信度
            if len(angles) > 1:
                second_variance = angles[1][1]
                confidence = max(0, 1.0 - best_variance / (second_variance + 1e-6))
            else:
                confidence = 0.5
            
            if abs(best_angle) > 1.0 and confidence > 0.3:
                corrected_image = self._rotate_image(image, -best_angle)
                return corrected_image, best_angle, {
                    "method": "key_alignment",
                    "confidence": confidence,
                    "key_count": len(key_contours),
                    "variance": best_variance
                }
            
            return None
            
        except Exception as e:
            self.logger.error(f"按键排列检测失败: {e}")
            return None
    
    def _select_best_correction(self, results: List[Optional[Tuple[npt.NDArray, float, dict]]]) -> Optional[Tuple[npt.NDArray, float, dict]]:
        """选择最佳的校正结果"""
        valid_results = [r for r in results if r is not None]
        
        if not valid_results:
            return None
        
        if len(valid_results) == 1:
            return valid_results[0]
        
        # 根据置信度和角度大小选择最佳结果
        best_result = None
        best_score = -1
        
        for result in valid_results:
            corrected_image, angle, info = result
            confidence = info.get('confidence', 0.0)
            angle_magnitude = abs(angle)
            
            # 综合评分：置信度 + 角度合理性
            score = confidence * 0.7 + min(angle_magnitude / 30.0, 1.0) * 0.3
            
            if score > best_score:
                best_score = score
                best_result = result
        
        return best_result
    
    def _rotate_image(self, image: npt.NDArray, angle: float) -> npt.NDArray:
        """旋转图像"""
        h, w = image.shape[:2]
        center = (w // 2, h // 2)
        
        # 计算旋转矩阵
        rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
        
        # 计算旋转后的图像尺寸
        cos_a = abs(rotation_matrix[0, 0])
        sin_a = abs(rotation_matrix[0, 1])
        new_w = int(h * sin_a + w * cos_a)
        new_h = int(h * cos_a + w * sin_a)
        
        # 调整旋转中心
        rotation_matrix[0, 2] += (new_w - w) / 2
        rotation_matrix[1, 2] += (new_h - h) / 2
        
        # 执行旋转
        rotated = cv2.warpAffine(image, rotation_matrix, (new_w, new_h), 
                                flags=cv2.INTER_CUBIC, 
                                borderMode=cv2.BORDER_REPLICATE)
        
        return rotated
    
    def visualize_rotation_correction(self, original: npt.NDArray, corrected: npt.NDArray, 
                                    angle: float, info: dict) -> npt.NDArray:
        """可视化旋转校正结果"""
        # 创建对比图像
        h1, w1 = original.shape[:2]
        h2, w2 = corrected.shape[:2]
        
        # 调整图像大小使其一致
        max_h = max(h1, h2)
        max_w = max(w1, w2)
        
        # 调整原始图像
        if h1 != max_h or w1 != max_w:
            original_resized = cv2.resize(original, (max_w, max_h))
        else:
            original_resized = original.copy()
        
        # 调整校正图像
        if h2 != max_h or w2 != max_w:
            corrected_resized = cv2.resize(corrected, (max_w, max_h))
        else:
            corrected_resized = corrected.copy()
        
        # 创建对比图像
        comparison = np.hstack([original_resized, corrected_resized])
        
        # 添加文字说明
        font = cv2.FONT_HERSHEY_SIMPLEX
        cv2.putText(comparison, f"Original (Rotated {angle:.1f}deg)", 
                   (10, 30), font, 0.8, (0, 255, 0), 2)
        cv2.putText(comparison, f"Corrected (Method: {info.get('method', 'unknown')})", 
                   (max_w + 10, 30), font, 0.8, (0, 255, 0), 2)
        cv2.putText(comparison, f"Confidence: {info.get('confidence', 0):.2f}", 
                   (max_w + 10, 60), font, 0.6, (255, 255, 0), 2)
        
        return comparison
